from sqlmodel import SQLModel
from core.Log import log_admin_event
from database.mysql import AsyncSessionLocal
from schemas.common import ListResponse, ListRequest, Filter
from fastapi import HTTPException
from sqlalchemy import text, func, delete
from sqlalchemy.orm import selectinload, joinedload, load_only
from sqlalchemy.future import select
from typing import List, Optional, Dict
import pandas as pd
from io import StringIO
import time

from server import util

class BaseServer:
    def __init__(self, model: SQLModel):
        self.model = model

    async def create(self, data: dict):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    obj = self.model(**data)
                    session.add(obj)
                    await session.commit()
                    return obj
            except Exception as e:
                await session.rollback()
                log_admin_event(f"create {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"create {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def update(self, data: dict):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    obj = await session.get(self.model, data["id"])
                    for key, value in data.items():
                        setattr(obj, key, value)
                    await session.commit()
                    return obj
            except Exception as e:
                await session.rollback()
                log_admin_event(f"update {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"update {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def delete(self, id: int):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    obj = await session.get(self.model, id)
                    await session.delete(obj)
                    await session.commit()
                    return True
            except Exception as e:
                await session.rollback()
                log_admin_event(f"delete {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"delete {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def deleteAll(self, ids: List[int]):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    for id in ids:
                        obj = await session.get(self.model, id)
                        await session.delete(obj)
                    await session.commit()
                    return True
            except Exception as e:
                await session.rollback()
                log_admin_event(f"deleteAll {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"deleteAll {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def get(self, id: int, fields: Optional[List[str]] = None, relations: Optional[Dict[str, List[str]]] = None):
        async with AsyncSessionLocal() as session:
            try:
                # 构建查询选项
                options = []
                
                # 处理关联模型的异步加载
                if relations:
                    for relation, rel_fields in relations.items():
                        rel_attr = getattr(self.model, relation)
                        # 使用 selectinload 进行异步加载关联模型
                        options.append(selectinload(rel_attr).load_only(*[getattr(rel_attr.mapper.class_, f) for f in rel_fields]))

                # 构建查询语句
                statement = select(self.model).options(*options).where(self.model.id == id)
                result = await session.execute(statement)
                result = result.first()

                if result:
                    result_obj = result[0]
                    # 如果指定了 fields，则只提取特定字段，否则返回整个对象
                    result_dict = result_obj.dict() if not fields else {field: getattr(result_obj, field) for field in fields}
                    
                    # 添加关联模型到返回结果中
                    if relations:
                        for relation in relations.keys():
                            result_dict[relation] = getattr(result[0], relation)

                    # 处理枚举和时间戳字段
                    # result_dict = util.process_enums(result_dict, result_obj)
                    # result_dict = util.process_timestamps(result_dict)
                    
                    return result_dict
                else:
                    raise HTTPException(404, {"msg": f"{self.model.__name__} not found", "detail": f"id: {id}"})
            except Exception as e:
                log_admin_event(f"get {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"get {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)
            
    async def get_one(self, where: List[Filter], fields: Optional[List[str]] = None, relations: Optional[Dict[str, List[str]]] = None):
        async with AsyncSessionLocal() as session:
            try:
                query = await self.build_query(where)

                options = []
                if fields:
                    options.append(load_only(*[getattr(self.model, field) for field in fields]))

                if relations:
                    for relation, rel_fields in relations.items():
                        rel_attr = getattr(self.model, relation)
                        options.append(joinedload(rel_attr).load_only(*[getattr(rel_attr.mapper.class_, f) for f in rel_fields]))

                query = query.options(*options)
                result = await session.execute(query)
                result = result.first()

                if result:
                    result_obj = result[0]
                    result_dict = result_obj.dict() if not fields else {field: getattr(result_obj, field) for field in fields}

                    if relations:
                        for relation in relations.keys():
                            result_dict[relation] = getattr(result[0], relation)

                    # result_dict = util.process_enums(result_dict, result_obj)
                    # result_dict = util.process_timestamps(result_dict)
                    return result_dict
                else:
                    raise HTTPException(404, {"msg": f"{self.model.__name__} not found", "detail": f"where: {where}"})
            except Exception as e:
                log_admin_event(f"get_one {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"get_one {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def list(self, request: ListRequest, fields: Optional[List[str]] = None, relations: Optional[Dict[str, List[str]]] = None):
        async with AsyncSessionLocal() as session:
            try:
                query = await self.build_query(request.filter)
                options = []
                if fields:
                    options.append(load_only(*[getattr(self.model, field) for field in fields]))

                if relations:
                    for relation, rel_fields in relations.items():
                        rel_attr = getattr(self.model, relation)
                        options.append(joinedload(rel_attr).load_only(*[getattr(rel_attr.mapper.class_, f) for f in rel_fields]))

                query = query.options(*options).order_by(text(f"{request.sort} {request.order}"))
                statement = query.limit(request.limit).offset((request.page - 1) * request.limit)
                results = await session.execute(statement)
                results = results.scalars().all()

                count_query = await self.build_query(request.filter)
                count_query = select(func.count()).select_from(count_query.subquery())
                total = await session.execute(count_query)
                total = total.scalar()

                result_list = []
                for result in results:
                    result_dict = result.dict() if not fields else {field: getattr(result, field) for field in fields}
                    if relations:
                        for relation in relations.keys():
                            result_dict[relation] = getattr(result, relation)

                    # result_dict = util.process_enums(result_dict, result)
                    result_list.append(result_dict)

                # result_list = util.process_timestamps(result_list)

                response = ListResponse(
                    list=result_list,
                    total=total,
                    pageSize=request.limit,
                    currentPage=request.page
                )
                return response
            except Exception as e:
                log_admin_event(f"list {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"list {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def all(self, filters: List[Filter] = [], fields: Optional[List[str]] = None, relations: Optional[Dict[str, List[str]]] = None):
        async with AsyncSessionLocal() as session:
            try:
                query = await self.build_query(filters)

                options = []
                if fields:
                    options.append(load_only(*[getattr(self.model, field) for field in fields]))

                if relations:
                    for relation, rel_fields in relations.items():
                        rel_attr = getattr(self.model, relation)
                        options.append(joinedload(rel_attr).load_only(*[getattr(rel_attr.mapper.class_, f) for f in rel_fields]))

                query = query.options(*options)
                results = await session.execute(query)
                results = results.scalars().all()

                result_list = []
                for result in results:
                    result_dict = result.dict() if not fields else {field: getattr(result, field) for field in fields}
                    if relations:
                        for relation in relations.keys():
                            result_dict[relation] = getattr(result, relation)

                    # result_dict = util.process_enums(result_dict, result)
                    result_list.append(result_dict)

                # result_list = util.process_timestamps(result_list)
                return result_list
            except Exception as e:
                log_admin_event(f"all {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"all {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def batch_delete(self, ids: List[int]):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    for id in ids:
                        obj = await session.get(self.model, id)
                        await session.delete(obj)
                    await session.commit()
                    return True
            except Exception as e:
                await session.rollback()
                log_admin_event(f"batch delete {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"batch delete {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def clear(self):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    await session.execute(delete(self.model))
                    await session.commit()
                    return True
            except Exception as e:
                await session.rollback()
                log_admin_event(f"clear {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"clear {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def change_status(self, id: int, status: int):
        async with AsyncSessionLocal() as session:
            try:
                async with session.begin():
                    obj = await session.get(self.model, id)
                    setattr(obj, "status", status)
                    await session.commit()
                    return True
            except Exception as e:
                await session.rollback()
                log_admin_event(f"change status {self.model.__name__} error: {str(e)}")
                return_data = {"msg": f"change status {self.model.__name__} error", "detail": str(e)}
                raise HTTPException(502, return_data)

    async def export(self, filters: List[Filter] = [], fields: Optional[List[str]] = None, relations: Optional[Dict[str, List[str]]] = None):
        filename = f"{self.model.__name__}_{time.strftime('%Y%m%d%H%M%S')}.csv"
        headers = {
            'Content-Disposition': 'attachment; filename=' + filename,
            'Content-Type': 'text/csv; charset=utf-8',
            'Access-Control-Expose-Headers': 'Content-Disposition'
        }

        data = await self.all(filters=filters, fields=fields, relations=relations)
        return data, headers

    async def generate_csv(self, data: List[Dict[str, str]]):
        buffer = StringIO()
        chunk_size = 10000
        for i in range(0, len(data), chunk_size):
            chunk = data[i:i + chunk_size]
            df = pd.DataFrame(chunk)
            df.to_csv(buffer, index=False, header=(i == 0))
            buffer.seek(0)
            yield buffer.read()
            buffer.truncate(0)
            buffer.seek(0)

    async def build_query(self, where: List[Filter]):
        query = select(self.model)
        for condition in where:
            field = condition.field
            operator = condition.operator
            value = condition.value
            column = getattr(self.model, field)

            if operator == "like":
                query = query.filter(column.like(f"%{value}%"))
            elif operator == "=":
                query = query.filter(column == value)
            elif operator == "!=":
                query = query.filter(column != value)
            elif operator == ">":
                query = query.filter(column > value)
            elif operator == "<":
                query = query.filter(column < value)
            elif operator == ">=":
                query = query.filter(column >= value)
            elif operator == "<=":
                query = query.filter(column <= value)
            elif operator == "between" and isinstance(value, list) and len(value) == 2:
                start_date, end_date = value
                query = query.filter(column.between(start_date, end_date))
            elif operator == "in" and isinstance(value, list):
                query = query.filter(column.in_(value))
        return query
